import json
from typing import List,Dict
import redis
from datetime import datetime

from click import prompt
from langchain_community.vectorstores import Milvus
from openai import vector_stores
from pymilvus import MilvusClient

from core.config import settings
from core.llm import get_default_llm,get_embedding_llm
from models.json_response  import JsonData
import logging
from langchain.retrievers.multi_query import MultiQueryRetriever

logger = logging.getLogger(__name__)
class ChatService:
    def __init__(self):
        self.milvus_client = MilvusClient(settings.MILVUS_URL)
        self.redis_client = redis.Redis(
            host=settings.REDIS_HOST, 
            port=settings.REDIS_PORT, 
            db=settings.REDIS_DB, 
            password=settings.REDIS_PASSWORD)
        
        self.llm = get_default_llm()

        self.embedding_llm = get_embedding_llm()
        
    def _get_chat_key(self,account_id:str)->str:
        return f"chat_history:{account_id}"
    
    def save_chat_history(self,account_id:str,message:List[Dict[str,str]])->None:
        key = self._get_chat_key(account_id)
        self.redis_client.lpush(key,json.dumps(message))
        
        
    def get_chat_history(self,account_id:str)->List[Dict[str,str]]:
        key = self._get_chat_key(account_id)
        messages = self.redis_client.lrange(key,0,-1)
        if messages:
            return [json.loads(message) for message in messages]
        return []
    
    def add_message(self,account_id:str,role:str,content:str)->None:
        message = self.get_chat_history(account_id)
        message.append({"role":role,"content":content,"timestamp":datetime.now().isoformat()})
        self.save_chat_history(account_id,message)
      
    def save_chat_message(self, account_id:str, user_message:str, assistant_message:str):
        self.add_message(account_id,"user",user_message)
        self.add_message(account_id,"assistant",assistant_message)
      
    def clear_chat_history(self,account_id:save_chat_history)->None:
        key = self._get_chat_key(account_id)
        self.redis_client.delete(key)
        
        
    def get_summary_key(self,account_id:str)->str:
        return f"summary:{account_id}"

    def get_information(self, message:str, course_id:str):
        #TODO 后续用多查询检索增强进行优化
        embedding = self.embedding_llm.embed_query(message)
        result = self.milvus_client.search(
            collection_name="teach_platform",
            data=[embedding],
            output_fields=["content"],
            filter=f"course_id=='{course_id}'",
        )
        return result
        
    async def generate_summary(self,account_id:str)->str:
        try:
            messages = self.get_chat_history(account_id)
            if not messages:
                return ""
            prompt = f"""请根据以下对话生成一个简要的核心摘要，突出话题和关键信息：
            
            {json.dumps(messages,ensure_ascii=False,indent=2)}
            
            摘要要求：
            1.突出对话的主要内容和关键信息
            2.使用第三人称描述，提取重要数据/时间节点/待办事项
            3.保留原始对话中的重要细节
            4.确保包含最新的对话内容
            
            """ 
            
            response = await self.llm.ainvoke(prompt)
            new_summary = response.content
            
            summary_key =  self.get_summary_key(account_id)
            old_summary = self.redis_client.get(summary_key)
            
            if old_summary:
                
                merged_prompt = f"""请将以下两个摘要合并成新的摘要：
                旧摘要：{old_summary}
                新摘要：{new_summary}
                
                和并要求：
                1.保留两个摘要中的重要信息
                2.突出对话的主要话题和关键信息
                3.使用第三人称描述，提取重要数据/时间节点/待办事项
                4.保留原始对话中的重要细节
                5.确保包含最新的对话内容
                """
                merged_prompt = await self.llm.ainvoke(prompt)
                final_summary = merged_prompt.content
            else:
                final_summary = new_summary
            
            self.redis_client.set(summary_key, final_summary)
        except Exception as e:
            logger.error("生成摘要失败{e}")
